package antonst.map;

import java.util.Comparator;
import java.util.Iterator;
import java.util.NoSuchElementException;

public class MyTreeMap<K, V> implements MyMap<K, V> {

	private final Comparator<K> comparator;
	private int size = 0;
	private SimpleEntry<K, V> root = null;

	public MyTreeMap() {
		comparator = null;
	}

	public MyTreeMap(Comparator<K> comparator) {
		this.comparator = comparator;
	}

	@Override
	public void clear() {
		root = null;
		size = 0;
	}

	@Override
	public boolean containsKey(Object key) {
		K keyTyped = (K) key;
		return findEntryWithKey(keyTyped) != null;
	}

	@Override
	public boolean containsValue(Object value) {
		return inorderTreeWalk(value, root);
	}
	
	@Override
	public V get(Object key) {
		K keyTyped = (K) key;
		SimpleEntry<K, V> entry = findEntryWithKey(keyTyped);
		if (entry != null) {
			return entry.getValue();
		} else {
			return null;
		}	
	}

	@Override
	public boolean isEmpty() {
		return size == 0;
	}

	@Override
	public V put(K key, V value) {
		SimpleEntry<K, V> alreadyExist = findEntryWithKey(key);
		if (alreadyExist != null) {
			alreadyExist.setValue(value);
		} else {
			SimpleEntry<K, V> insertedEntry = insertInTree(key, value);
			insertFixUp(insertedEntry);
			size = size + 1;
		}
		return value;
	}

	@Override
	public V remove(Object key) {
		K keyTyped = (K) key;
		SimpleEntry<K, V> entryToDelete = findEntryWithKey(keyTyped);
		V value = entryToDelete.getValue();
		if (entryToDelete != null) {
			SimpleEntry<K, V> entryForFixUp = rbDelete(entryToDelete);
			if (entryForFixUp != null)
				rbDeleteFixUp(entryForFixUp);
			size = size - 1;
			return value;
		} else {
			return null;
		}
	}

	@Override
	public int size() {
		return size;
	}

	@Override
	public Iterator<Entry<K, V>> entryIterator() {
		return new MyIterator();
	}

	private SimpleEntry<K, V> insertInTree(K key, V value) {
		SimpleEntry<K, V> parentEntry = null;
		SimpleEntry<K, V> currentEntry = root;
		while (currentEntry != null) {
			parentEntry = currentEntry;
			if (compareTwoKeys(key, currentEntry.getKey()) < 0) {
				currentEntry = currentEntry.left;
			} else {
				currentEntry = currentEntry.right;
			}
		}
		SimpleEntry<K, V> insertedEntry = new SimpleEntry<>(key, value);
		insertedEntry.parent = parentEntry;
		if (parentEntry == null) {
			root = insertedEntry;
		} else {
			if (compareTwoKeys(key, parentEntry.getKey()) < 0) {
				parentEntry.left = insertedEntry;
			} else {
				parentEntry.right = insertedEntry;
			}
		}
		return insertedEntry;
	}

	private int compareTwoKeys(K key1, K key2) {
		if (key1 == null && key2 == null) { 
			return 0;
		} else if (key1 == null) {
			return -1;
		} else if (key2 == null) {
			return 1;
		}
		if (comparator == null) {
			Comparable<K> key1Comparable = (Comparable<K>) key1;
			return key1Comparable.compareTo(key2);
		} else {
			return comparator.compare(key1, key2);
		}
	}

	private SimpleEntry<K, V> findEntryWithKey(K key) {
		SimpleEntry<K, V> entry = root;
		while (entry != null) {
			if (compareTwoKeys(key, entry.getKey()) == 0) {
				return entry;	
			} else if (compareTwoKeys(key, entry.getKey()) < 0) {
				entry = entry.left;
			} else {
				entry = entry.right;
			}
		}
		return null;
	}

	private void insertFixUp(SimpleEntry<K, V> entry) {
		while (entry.parent != null && entry.parent.isRed()) {
			if (entry.parent == entry.parent.parent.left) {
				SimpleEntry<K, V> parentParentRight = entry.parent.parent.right;
				if (parentParentRight.isRed()) {
					entry.parent.repaintBlack();
					parentParentRight.repaintBlack();
					entry.parent.parent.repaintRed();
					entry = entry.parent.parent;
				} else if (entry == entry.parent.right) {
					entry = entry.parent;
					leftRotate(entry);
				} else {
					entry.parent.repaintBlack();
					entry.parent.parent.repaintRed();
					rightRotate(entry.parent.parent);
				}
			} else {
				SimpleEntry<K, V> parentParentLeft = entry.parent.parent.left;
				if (parentParentLeft != null && parentParentLeft.isRed()) {
					entry.parent.repaintBlack();
					parentParentLeft.repaintBlack();
					entry.parent.parent.repaintRed();
					entry = entry.parent.parent;
				} else if (entry == entry.parent.left) {
					entry = entry.parent;
					rightRotate(entry);
				} else {
					entry.parent.repaintBlack();
					entry.parent.parent.repaintRed();
					leftRotate(entry.parent.parent);
				}
			}
		}
		root.repaintBlack();
	}

	private void leftRotate(SimpleEntry<K, V> entry) {
		SimpleEntry<K, V> newTopEntry = entry.right;
		entry.right = newTopEntry.left;
		if (newTopEntry.left != null)
			newTopEntry.left.parent = entry;
		newTopEntry.parent = entry.parent;
		if (entry.parent == null) {
			root = newTopEntry;
		} else if (entry == entry.parent.left) {
			entry.parent.left = newTopEntry;
		} else {
			entry.parent.right = newTopEntry;
		}
		newTopEntry.left = entry;
		entry.parent = newTopEntry;
	}

	private void rightRotate(SimpleEntry<K, V> entry) {
		SimpleEntry<K, V> newTopEntry = entry.left;
		entry.left = newTopEntry.right;
		if (newTopEntry.right != null) {
			newTopEntry.right.parent = entry;
		}
		newTopEntry.parent = entry.parent;
		if (entry.parent == null) {
			root = newTopEntry;
		} else if (entry == entry.parent.left) {
			entry.parent.left = newTopEntry;
		} else {
			entry.parent.right = newTopEntry;
		}
	}

	private boolean inorderTreeWalk(Object value, SimpleEntry<K, V> entry) {
		if (entry == null) {
			return false;
		} else if (inorderTreeWalk(value, entry.left)) { 
			return true;
		} else if (entry.getValue().equals(value)) {
			return true;
		} else if (inorderTreeWalk(value, entry.right)) {
			return true;
		} else {
			return false;
		}
	}

	private SimpleEntry<K, V> getSuccessor(SimpleEntry<K, V> entry) {
		if (entry.right != null) {
			return treeMinimum(entry.right);
		}	
		SimpleEntry<K, V> parentEntry = entry.parent;
		while (parentEntry != null && entry == parentEntry.right) {
			entry = parentEntry;
			parentEntry = parentEntry.parent;
		}
		return parentEntry;
	}

	private SimpleEntry<K, V> treeMinimum(SimpleEntry<K, V> entry) {
		if (entry == null) return null;
		while (entry.left != null)	
			entry = entry.left;
		return entry;
	}

	private SimpleEntry<K, V> rbDelete(SimpleEntry<K, V> entry) {
		SimpleEntry<K, V> y;
		if (entry.left == null || entry.right == null) {
			y = entry;
		} else {
			y = getSuccessor(entry);
		}
		SimpleEntry<K, V> x;
		if (y.left != null) {
			x = y.left;
		} else {
			x = y.right;
		}
		if (x == null) 
			x = new SimpleEntry<>(null, null);
		x.parent = y.parent;
		if (y.parent == null) {
			root = x;
		} else if (y == y.parent.left) {
			y.parent.left = x;
		} else {
			y.parent.right = x;
		}
		if (y != entry) {
			entry.key = y.key;
			entry.value = y.value;
			entry.left = y.left;
			entry.right = y.right;
			entry.parent = y.parent;
			entry.color = y.color;
		}
		if (y.isBlack()) {
			return x;
		} else {
			return null;
		}
	}

	private void rbDeleteFixUp(SimpleEntry<K, V> entry) {
		while (entry != root && entry.isBlack()) {
			if (entry == entry.parent.left) {
				SimpleEntry<K, V> w = entry.parent.right;
				if (w.isRed()) {
					w.repaintBlack();
					entry.parent.repaintRed();
					leftRotate(entry.parent);
					w = entry.parent.right;
				}
				if (w.left.isBlack() && w.right.isBlack()) {
					w.repaintRed();
					entry = entry.parent;
				} else if (w.right.isBlack()) {
					w.left.repaintBlack();
					w.repaintRed();
					rightRotate(w);
					w = entry.parent.right;
				}
				w.color = entry.parent.color;
				entry.parent.repaintBlack();
				w.right.repaintBlack();
				leftRotate(entry.parent);
				entry = root;
			} else {
				SimpleEntry<K, V> w = entry.parent.left;
				if (w.isRed()) {
					w.repaintBlack();
					entry.parent.repaintRed();
					rightRotate(entry.parent);
					w = entry.parent.left;
				}
				if (w.left.isBlack() && w.right.isBlack()) {
					w.repaintRed();
					entry = entry.parent;
				} else if (w.left.isBlack()) {
					w.right.repaintBlack();
					w.repaintRed();
					leftRotate(w);
					w = entry.parent.left;
				}
				w.color = entry.parent.color;
				entry.parent.repaintBlack();
				w.left.repaintBlack();
				rightRotate(entry.parent);
				entry = root;
			}
		}
		entry.repaintBlack();
	}

	public static class SimpleEntry<K, V> implements MyMap.Entry<K, V> {

		public enum Color {

			BLACK, RED;
		}

		Color color = Color.RED;
		SimpleEntry<K, V> right;
		SimpleEntry<K, V> left;
		SimpleEntry<K, V> parent;
		private K key;
		private V value;

		public SimpleEntry(K key, V value) {
			this.key = key;
			this.value = value;
		}

		@Override
		public K getKey() {
			return key;
		}

		@Override
		public V getValue() {
			return value;
		}

		@Override
		public V setValue(V value) {
			this.value = value;
			return value;
		}

		boolean isRed() {
			return color == Color.RED;
		}

		boolean isBlack() {
			return color == Color.BLACK;
		}

		void repaintRed() {
			color = Color.RED;
		}

		void repaintBlack() {
			color = Color.BLACK;
		}

		@Override
		public boolean equals(Object obj) {
			if (obj == null) return false;
			if (obj instanceof SimpleEntry){
				SimpleEntry other = (SimpleEntry) obj; 
				if (key == null && other.key == null) {
					return value.equals(other.value);
				} else if (value == null && other.value == null) {
					return key.equals(other.key);
				} else {
					return key.equals(other.key) & value.equals(other.value);
				}
			}
			return false;
		} 

		@Override
		public int hashCode() {
			int keyHashCode = (key == null) ? 0 : key.hashCode();
			int valueHashCode = (value == null) ? 0 : value.hashCode();
			return keyHashCode ^ valueHashCode;
		}
	}

	private class MyIterator implements Iterator<Entry<K, V>> {

		private SimpleEntry<K, V> lastReturned = null;
		private SimpleEntry<K, V> next;
		
		public MyIterator() {
			next = treeMinimum(MyTreeMap.this.root);
		}
		
		@Override
		public boolean hasNext() {
			return next != null;
		}

		@Override
		public Entry<K, V> next() {
			if (next != null) {
				lastReturned = next;
				next = getSuccessor(next);
				return lastReturned;
			} else {
				throw new NoSuchElementException();
			}
		}

		@Override
		public void remove() {
			throw new UnsupportedOperationException("Not supported yet."); //To change body of generated methods, choose Tools | Templates.
		}
		
	}

}
